{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EHq8UWSjXSyz"
   },
   "source": [
    "<center>\n",
    "<h4>CDS 110, Lecture 4b</h4>\n",
    "<font color=blue><h1>LQR Tracking</h1></font>\n",
    "<h3>Richard M. Murray and Natalie Bernat, Winter 2024</h3>\n",
    "</center>\n",
    "\n",
    "[Open in Google Colab](https://colab.research.google.com/drive/1Q6hXokOO_e3-wl6_ghigpxGJRUrGcHp3)\n",
    "\n",
    "This example uses a linear system to show how to implement LQR based tracking and some of the tradeoffs between feedfoward and feedback.  Integral action is also implemented."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "try:\n",
    "  import control as ct\n",
    "  print(\"python-control\", ct.__version__)\n",
    "except ImportError:\n",
    "  !pip install control\n",
    "  import control as ct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "a23d6f89"
   },
   "source": [
    "# Part I: Second order linear system\n",
    "\n",
    "We'll use a simple linear system to illustrate the concepts:\n",
    "$$\n",
    "\\frac{dx}{dt} =\n",
    "\\begin{bmatrix}\n",
    "0 & 10 \\\\\n",
    "-1 & 0\n",
    "\\end{bmatrix}\n",
    "x +\n",
    "\\begin{bmatrix}\n",
    "0  \\\\\n",
    "1\n",
    "\\end{bmatrix}\n",
    "u,\n",
    "\\qquad\n",
    "y = \\begin{bmatrix} 1 & 1 \\end{bmatrix} x.\n",
    "$$\n",
    "\n",
    "<!-- This system corresponds to the linearized lateral dynamics of a vehicle driving down a road at 10 m/s. -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a simple linear system that we want to control\n",
    "A = np.array([[0, 10], [-1, 0]])\n",
    "B = np.array([[0], [1]])\n",
    "C = np.array([[1, 1]])\n",
    "sys = ct.ss(A, B, C, 0, name='sys')\n",
    "print(sys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ja1g1MlbieJy"
   },
   "source": [
    "## Linear quadratic regulator (LQR) design\n",
    "\n",
    "We'll design a controller of the form\n",
    "\n",
    "$$\n",
    "u=-Kx+k_rr\n",
    "$$\n",
    "\n",
    "- For the feedback control gain $K$, we'll use linear quadratic regulator theory.  We seek to find the control law that minimizes the cost function:\n",
    "\n",
    "  $$\n",
    "  J(x(\\cdot), u(\\cdot)) = \\int_0^\\infty x^T(\\tau) Q x(\\tau) + u^T(\\tau) R u(\\tau)\\, d\\tau\n",
    "  $$\n",
    "\n",
    "  The weighting matrices $Q\\succeq 0 \\in \\mathbb{R}^{n \\times n}$ and $R \\succ 0\\in \\mathbb{R}^{m \\times m}$ should be chosen based on the desired performance of the system (tradeoffs in state errors and input magnitudes).  See Example 3.5 in [Optimization Based Control (OBC)](https://fbswiki.org/wiki/index.php/Supplement:_Optimization-Based_Control) for a discussion of how to choose these weights.  For now, we just choose identity weights for all states and inputs.\n",
    "\n",
    "- For the feedforward control gain $k_r$, we derive the feedforward gain from an equilibrium point analysis:\n",
    "  $$\n",
    "  y_e = C(A-BK)^{-1}Bk_rr\n",
    "  \\qquad\\implies\\qquad k_r = \\frac{-1}{C(A-BK)^{-1}B}\n",
    "  $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct an LQR controller for the system\n",
    "Q = np.eye(sys.nstates)\n",
    "R = np.eye(sys.ninputs)\n",
    "K, _, _ = ct.lqr(sys, Q, R)\n",
    "print('K: '+str(K))\n",
    "\n",
    "# Set the feedforward gain to track the reference\n",
    "kr = (-1 / (C @ np.linalg.inv(A - B @ K) @ B))\n",
    "print('k_r: '+str(kr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "99f036ea"
   },
   "source": [
    "Now that we have our gains designed, we can simulate the closed loop system:\n",
    "$$\n",
    "\\frac{dx}{dt} = A_{cl}x + B_{cl} r,\n",
    "\\quad A_{cl} = A-BK,\n",
    "\\quad B_{cl} = Bk_r\n",
    "$$\n",
    "Notice that, with a state feedback controller, the new (closed loop) dynamics matrix absorbs the old (open loop) \"input\" $u$, and the new (closed loop) input is our reference signal $r$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a closed loop system\n",
    "A_cl = A - B @ K\n",
    "B_cl =  B * kr\n",
    "clsys = ct.ss(A_cl, B_cl, C, 0)\n",
    "print(clsys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "84422c3f"
   },
   "source": [
    "## System simulations\n",
    "\n",
    "### Baseline controller\n",
    "\n",
    "To see how the baseline controller performs, we ask it to track a constant reference $r = 2$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the step response with respect to the reference input\n",
    "r = 2\n",
    "Tf = 8\n",
    "tvec = np.linspace(0, Tf, 100)\n",
    "\n",
    "U = r * np.ones_like(tvec)\n",
    "time, output = ct.input_output_response(clsys, tvec, U)\n",
    "plt.plot(time, output)\n",
    "plt.plot([time[0], time[-1]], [r, r], '--');\n",
    "plt.legend(['y', 'r']);\n",
    "plt.ylabel(\"Output\")\n",
    "plt.xlabel(\"Time $t$ [sec]\")\n",
    "plt.title(\"Baseline controller step response\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ea2d1c59"
   },
   "source": [
    "Things to try:\n",
    "- set $k_r=0$\n",
    "- set $k_r \\neq \\frac{-1}{C(A-BK)^{-1}B}$\n",
    "- try different LQR weightings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "84ee7635"
   },
   "source": [
    "### Disturbance rejection\n",
    "\n",
    "To add an input disturbance to the system, we include a second open loop input:\n",
    "$$\n",
    "\\frac{dx}{dt} =\n",
    "\\begin{bmatrix}\n",
    "0 & 10 \\\\\n",
    "-1 & 0\n",
    "\\end{bmatrix}\n",
    "x +\n",
    "\\begin{bmatrix}\n",
    "0 & 0\\\\\n",
    "1 & 1\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "u\\\\\n",
    "d\n",
    "\\end{bmatrix},\n",
    "\\qquad\n",
    "y = \\begin{bmatrix} 1 & 1 \\end{bmatrix} x.\n",
    "$$\n",
    "\n",
    "Our closed loop system becomes:\n",
    "$$\n",
    "\\frac{dx}{dt} =\n",
    "\\begin{bmatrix}\n",
    "0 & 10 \\\\\n",
    "-1-K_{1} & 0-K_{2}\n",
    "\\end{bmatrix}\n",
    "x +\n",
    "\\begin{bmatrix}\n",
    "0 & 0\\\\\n",
    "k_r & 1\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "r\\\\\n",
    "d\n",
    "\\end{bmatrix},\n",
    "\\qquad\n",
    "y = \\begin{bmatrix} 1 & 1 \\end{bmatrix} x.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Resimulate with a disturbance input\n",
    "B_ext = np.hstack([B * kr, B])\n",
    "clsys = ct.ss(A - B @ K, B_ext, C, 0)\n",
    "\n",
    "# Construct the inputs for the augmented system\n",
    "delta = 0.5\n",
    "U = np.vstack([r * np.ones_like(tvec), delta * np.ones_like(tvec)])\n",
    "\n",
    "time, output = ct.input_output_response(clsys, tvec, U)\n",
    "\n",
    "plt.plot(time, output[0])\n",
    "plt.plot([time[0], time[-1]], [r, r], '--')\n",
    "plt.legend(['y', 'r']);\n",
    "plt.ylabel(\"Output\")\n",
    "plt.xlabel(\"Time $t$ [sec]\")\n",
    "plt.title(\"Baseline controller step response with disturbance\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Qis2PP3nd7ua"
   },
   "source": [
    "We see that this leads to steady state error, since the feedforward signal didn't include an offset for the disturbance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "84a9e61c"
   },
   "source": [
    "#### Integral feedback\n",
    "\n",
    "A standard approach to compensate for constant disturbances is to use integral feedback.  To do this, we have to keep track of the integral of the error\n",
    "\n",
    "$$z = \\int_0^\\tau (y - r)\\, d\\tau= \\int_0^\\tau (Cx - r)\\, d\\tau.$$\n",
    "\n",
    "We do this by creating an augmented system that includes the dynamics of the process ($dx/dt$) along with the dynamics of the integrator state ($dz/dt$):\n",
    "\n",
    "$$\n",
    "\\frac{d}{dt}\\begin{bmatrix}\n",
    "x \\\\\n",
    "z\n",
    "\\end{bmatrix} =\n",
    "\\begin{bmatrix}\n",
    "A & 0 \\\\\n",
    "C & 0\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "x \\\\\n",
    "z\n",
    "\\end{bmatrix} +\n",
    "\\begin{bmatrix}\n",
    "B\\\\\n",
    "0 \\\\\n",
    "\\end{bmatrix}\n",
    "u+\n",
    "\\begin{bmatrix}\n",
    "0\\\\\n",
    "-I \\\\\n",
    "\\end{bmatrix}\n",
    "r,\n",
    "\\qquad\n",
    "y = \\begin{bmatrix} C \\\\ 0 \\end{bmatrix} \\begin{bmatrix}\n",
    "x \\\\\n",
    "z\n",
    "\\end{bmatrix}.\n",
    "$$\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define an augmented state space for use with LQR\n",
    "A_aug = np.block([[sys.A, np.zeros((sys.nstates, 1))], [C, 0] ])\n",
    "B_aug = np.vstack([sys.B, 0])\n",
    "print(\"A =\", A_aug, \"\\nB =\", B_aug)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "463d9b85"
   },
   "source": [
    "\n",
    "Our controller then takes the form:\n",
    "\n",
    "$$\n",
    "\\begin{aligned}\n",
    "u &= - Kx - k_\\text{i} \\int_0^\\tau (y - r)\\, d\\tau+k_rr \\\\\n",
    " &= - (Kx + k_\\text{i}z)+k_rr .\n",
    "\\end{aligned}\n",
    "$$\n",
    "\n",
    "This results in the closed loop system:\n",
    "$$\n",
    "\\frac{dx}{dt} =\n",
    "\\begin{bmatrix}\n",
    "A-BK & -Bk_i \\\\\n",
    "C & 0\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "x \\\\\n",
    "z\n",
    "\\end{bmatrix} +\n",
    "\\begin{bmatrix}\n",
    "Bk_r\\\\\n",
    "-I \\\\\n",
    "\\end{bmatrix}\n",
    "r,\n",
    "\\qquad\n",
    "y = \\begin{bmatrix} C \\\\ 0 \\end{bmatrix} \\begin{bmatrix}\n",
    "x \\\\\n",
    "z\n",
    "\\end{bmatrix}.\n",
    "$$\n",
    "\n",
    "Since z is part of the augmented state space, we can generate an LQR controller for the augmented system to find both the usual gain $K$ and the integral gain $k_i$:\n",
    "$$\n",
    "\\bar{K} = \\begin{bmatrix} K& k_i\\end{bmatrix}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an LQR controller for the augmented system\n",
    "K_aug, _, _ = ct.lqr(A_aug, B_aug, np.diag([1, 1, 1]), np.eye(sys.ninputs))\n",
    "print('K_aug: '+str(K_aug))\n",
    "\n",
    "K = K_aug[:, 0:2]\n",
    "ki = K_aug[:, 2]\n",
    "kr = -1 / (C @ np.linalg.inv(A - B * K) @ B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "19bb6592"
   },
   "source": [
    "<!-- We can think about this gain as `K_aug = [K, ki]` and the resulting contoller becomes -->\n",
    "\n",
    "\n",
    "Notice that the value of $K$ changed, so we needed to recompute $k_r$ too."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zHlf8zoHoqvF"
   },
   "source": [
    "To run simulations, we return to our system augmented with a disturbance, but we expand the outputs available to the controller:\n",
    "\n",
    "$$\n",
    "\\frac{dx}{dt} =\n",
    "\\begin{bmatrix}\n",
    "0 & 10 \\\\\n",
    "-1 & 0\n",
    "\\end{bmatrix}\n",
    "x +\n",
    "\\begin{bmatrix}\n",
    "0 & 0\\\\\n",
    "1 & 1\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "u\\\\\n",
    "d\n",
    "\\end{bmatrix},\n",
    "$$\n",
    "\n",
    "$$\n",
    "\\bar{y} = \\begin{bmatrix} 1 & 0 & 1 \\\\ 0 & 1 & 1 \\end{bmatrix}^T x = \\begin{bmatrix} x_1 & x_2 & y \\end{bmatrix}  .\n",
    "$$\n",
    "\n",
    "The controller then constructs its internal state $z$ out of $x$ and $r$.\n",
    "\n",
    "<!-- $$\n",
    "\\frac{dx}{dt} =\n",
    "\\begin{bmatrix}\n",
    "A-BK & -Bk_i \\\\\n",
    "C & 0\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "x \\\\\n",
    "z\n",
    "\\end{bmatrix} +\n",
    "\\begin{bmatrix}\n",
    "Bk_r & B\\\\\n",
    "-I & 0 \\\\\n",
    "\\end{bmatrix}\n",
    "\\begin{bmatrix}\n",
    "r \\\\\n",
    "d\n",
    "\\end{bmatrix},\n",
    "\\qquad\n",
    "y = \\begin{bmatrix} C \\\\ 0 \\end{bmatrix} \\begin{bmatrix}\n",
    "x \\\\\n",
    "z\n",
    "\\end{bmatrix}.\n",
    "$$ -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Construct a system with disturbance inputs, and full outputs (for the controller)\n",
    "A_integral = sys.A\n",
    "B_integral = np.hstack([sys.B, sys.B])\n",
    "C_integral = [[1, 0], [0, 1], [1, 1]] # outputs for the controller: x1, x2, y\n",
    "sys_integral = ct.ss(\n",
    "    A_integral, B_integral, C_integral, 0,\n",
    "    inputs=['u', 'd'],\n",
    "    outputs=['x1', 'x2', 'y']\n",
    ")\n",
    "print(sys_integral)\n",
    "\n",
    "# Construct an LQR+integral controller for the system with an internal state z\n",
    "A_ctrl = [[0]]\n",
    "B_ctrl = [[1, 1, -1]] # z_dot=Cx-r\n",
    "C_ctrl = -ki #-ki*z\n",
    "D_ctrl = np.hstack([-K, kr]) #-K*x + kr*r\n",
    "ctrl_integral=ct.ss(\n",
    "    A_ctrl, B_ctrl, C_ctrl, D_ctrl, # u = -ki*z - K*x + kr*r\n",
    "    inputs=['x1', 'x2', 'r'],    # system outputs + reference\n",
    "    outputs=['u'],               # controller action\n",
    ")\n",
    "print(ctrl_integral)\n",
    "\n",
    "# Create the closed loop system\n",
    "clsys_integral = ct.interconnect([sys_integral, ctrl_integral], inputs=['r', 'd'], outputs=['y'])\n",
    "print(clsys_integral)\n",
    "\n",
    "# Resimulate with a disturbance input\n",
    "delta = 0.5\n",
    "U = np.vstack([r * np.ones_like(tvec), delta * np.ones_like(tvec)])\n",
    "time, output, states = ct.input_output_response(clsys_integral, tvec, U, return_x=True)\n",
    "plt.plot(time, output[0])\n",
    "plt.plot([time[0], time[-1]], [r, r], '--')\n",
    "plt.plot(time, states[2])\n",
    "plt.legend(['y', 'r', 'z']);\n",
    "plt.ylabel(\"Output\")\n",
    "plt.xlabel(\"Time $t$ [sec]\")\n",
    "plt.title(\"LQR+integral controller step response with disturbance\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M9nXbITrhYg7"
   },
   "source": [
    "Notice that the steady state value of $z=\\int(y-r)$ is not zero, but rather settles to whatever value makes $y-r$ zero!\n",
    "<!-- (If there's time: try setting ud=0 again to see what changes for the integral controller) -->"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f8bfc15c"
   },
   "source": [
    "# Part II: PVTOL Linear Quadratic Regulator Example\n",
    "\n",
    "Natalie Bernat, 26 Apr 2024 <br>\n",
    "Richard M. Murray, 25 Jan 2022\n",
    "\n",
    "This notebook contains an example of LQR control applied to the PVTOL system.  It demonstrates how to construct an LQR controller by linearizing the system, and provides an alternate view of the feedforward component of the controller."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "77e2ed47"
   },
   "source": [
    "## System description\n",
    "\n",
    "We use the PVTOL dynamics from [Feedback Systems (FBS2e)](https://fbswiki.org/wiki/index.php/Feedback_Systems:_An_Introduction_for_Scientists_and_Engineers), which can be found in Example 3.12}\n",
    "\n",
    "<table>\n",
    "<tr>\n",
    "    <td width=\"30%\"><img src=\"https://fbswiki.org/wiki/images/7/76/Pvtol.png\" width=240></td>\n",
    "    <td width=\"30%\">\n",
    "$$\n",
    "\\begin{aligned}\n",
    "  m \\ddot x &= F_1 \\cos\\theta - F_2 \\sin\\theta - c \\dot x, \\\\\n",
    "  m \\ddot y &= F_1 \\sin\\theta + F_2 \\cos\\theta - m g - c \\dot y, \\\\\n",
    "  J \\ddot \\theta &= r F_1.\n",
    "\\end{aligned}\n",
    "$$\n",
    "    </td>\n",
    "    <td width=\"30%\">\n",
    "$$\n",
    "\\frac{dz}{dt} =\n",
    "\\begin{bmatrix}\n",
    "z_4 \\\\\n",
    "z_5 \\\\\n",
    "z_6 \\\\\n",
    "-\\frac{c}{m}z_4 \\\\\n",
    "-g-\\frac{c}{m}z_5 \\\\\n",
    "0\n",
    "\\end{bmatrix} +\n",
    "\\begin{bmatrix}\n",
    "0 \\\\\n",
    "0 \\\\\n",
    "0 \\\\\n",
    "\\frac{F_1}{m}cos\\theta -\\frac{F_2}{m}sin\\theta  \\\\\n",
    "\\frac{F_1}{m}sin\\theta +\\frac{F_2}{m}cos\\theta \\\\\n",
    "-\\frac{r}{J}F_1\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "    </td>\n",
    "</tr>\n",
    "</table>\n",
    "\n",
    "The state space variables for this system are:\n",
    "\n",
    "$z=(x,y,\\theta, \\dot x,\\dot y,\\dot  \\theta), \\quad u=(F_1,F_2)$\n",
    "\n",
    "Notice that the x and y positions ($z_1$ and $z_2$) do not actually appear in the dynamics-- this makes sense, since the aircraft should hypothetically fly the same way no matter where in the air it is (neglecting effects near the ground)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# PVTOL dynamics\n",
    "def pvtol_update(t, x, u, params):\n",
    "    from math import cos, sin\n",
    "    \n",
    "    # Get the parameter values\n",
    "    m, J, r, g, c = map(params.get, ['m', 'J', 'r', 'g', 'c'])\n",
    "\n",
    "    # Get the inputs and states\n",
    "    x, y, theta, xdot, ydot, thetadot = x\n",
    "    F1, F2 = u\n",
    "\n",
    "    # Constrain the inputs\n",
    "    F2 = np.clip(F2, 0, 1.5 * m * g)\n",
    "    F1 = np.clip(F1, -0.1 * F2, 0.1 * F2)\n",
    "\n",
    "    # Dynamics\n",
    "    xddot = (F1 * cos(theta) - F2 * sin(theta) - c * xdot) / m\n",
    "    yddot = (F1 * sin(theta) + F2 * cos(theta) - m * g - c * ydot) / m\n",
    "    thddot = (r * F1) / J\n",
    "\n",
    "    return np.array([xdot, ydot, thetadot, xddot, yddot, thddot])\n",
    "\n",
    "def pvtol_output(t, x, u, params):\n",
    "    return x\n",
    "\n",
    "pvtol = ct.nlsys(\n",
    "    pvtol_update, pvtol_output, name='pvtol',\n",
    "    states = [f'x{i}' for i in range(6)],\n",
    "    inputs = ['F1', 'F2'],\n",
    "    outputs=[f'x{i}' for i in range(6)],\n",
    "    # outputs = ['x', 'y', 'theta', 'xdot', 'ydot', 'thdot'],\n",
    "    params = {\n",
    "        'm': 4.,                # mass of aircraft\n",
    "        'J': 0.0475,            # inertia around pitch axis\n",
    "        'r': 0.25,              # distance to center of force\n",
    "        'g': 9.8,               # gravitational constant\n",
    "        'c': 0.05,              # damping factor (estimated)\n",
    "    }\n",
    ")\n",
    "\n",
    "print(pvtol)\n",
    "print(pvtol.params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YZiISLS-qMS_"
   },
   "source": [
    "Next, we'll linearize the system around the equilibrium points. As discussed in FBS2e (example 7.9), the linearization around this equilibrium point has the form:\n",
    "$$\n",
    "A =\n",
    "\\begin{bmatrix}\n",
    "0 & 0 & 0 & 1 & 0 & 0\\\\\n",
    "0 & 0 & 0 & 0 & 1 & 0 \\\\\n",
    "0 & 0 & 0 & 0 & 0 & 1 \\\\\n",
    "0 & 0 & -g & -c/m & 0 & 0 \\\\\n",
    "0 & 0 & 0 & 0 & -c/m & 0 \\\\\n",
    "0 & 0 & 0 & 0 & 0 & 0\n",
    "\\end{bmatrix}\n",
    ", \\quad B=\n",
    "\\begin{bmatrix}\n",
    "0 & 0 \\\\\n",
    "0 & 0 \\\\\n",
    "0 & 0 \\\\\n",
    "1/m & 0 \\\\\n",
    "0 & 1/m \\\\\n",
    "r/J & 0\n",
    "\\end{bmatrix}\n",
    ".\n",
    "$$\n",
    "(note that here $r$ is a system parameter, not the same as the reference $r$ we've been using elsewhere in this notebook)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To compute this linearization in python-control, we start by computing the equilibrium point.  We do this using the `find_eqpt` function, which can be used to find equilibrium points satisfying varioius conditions.  For this system, we wish to find the state $x_\\text{e}$ and input $u_\\text{e}$ that holds the $x, y$ position of the aircraft at the point $(0, 0)$.  The `find_eqpt` function performs a numerical optimization to find the values of $x_\\text{e}$ and $u_\\text{e}$ corresponding to an equilibrium point with the desired values for the outputs.  We pass the function initial guesses for the state and input as well the values of the output and the indices of the output that we wish to constrain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the equilibrium point corresponding to hover\n",
    "xeq, ueq = ct.find_eqpt(pvtol, np.zeros(6), np.zeros(2), y0=np.zeros(6), iy=[0, 1])\n",
    "print(f\"{xeq=}, {ueq=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using these values, we compute the linearization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "linsys = pvtol.linearize(xeq, ueq)\n",
    "print(linsys)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7cb8840b"
   },
   "source": [
    "## Linear quadratic regulator (LQR) design\n",
    "\n",
    "Now that we have a linearized model of the system, we can compute a controller using linear quadratic regulator theory. We wish to minimize the following cost function\n",
    "\n",
    "$$\n",
    "J(\\phi(\\cdot), \\nu(\\cdot)) = \\int_0^\\infty \\phi^T(\\tau) Q \\phi(\\tau) + \\nu^T(\\tau) R \\nu(\\tau)\\, d\\tau,\n",
    "$$\n",
    "\n",
    "where we have changed to our linearized coordinates:\n",
    "\n",
    "$$\\phi=z-z_e, \\quad \\nu = u-u_e$$\n",
    "\n",
    "Using the standard approach for finding K, we obtain a feedback controller for the system:\n",
    "$$\\nu=-K\\phi$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start with a diagonal weighting\n",
    "Q1 = np.diag([1, 1, 1, 1, 1, 1])\n",
    "R1 = np.diag([1, 1])\n",
    "K, X, E = ct.lqr(linsys, Q1, R1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "863d07de"
   },
   "source": [
    "To create a controller for the system, we have to apply a control signal $u$, so we change back from the relative coordinates to the absolute coordinates:\n",
    "\n",
    "$$u=u_e - K(z - z_e)$$\n",
    "\n",
    "Notice that, since $(Kz_e+u_e)$ is completely determined by (user-defined) inputs to the system, this term is a type of feedforward control signal.\n",
    "\n",
    "To create a controller for the system, we can use the function  [`create_statefbk_iosystem()`](https://python-control.readthedocs.io/en/latest/generated/control.create_statefbk_iosystem.html), which creates an I/O system that takes in a desired trajectory $(x_\\text{d}, u_\\text{d})$ and the current state $x$ and generates a control law of the form:\n",
    "\n",
    "$$\n",
    "u = u_\\text{d} - K (x - x_\\text{d})\n",
    "$$\n",
    "\n",
    "Note that this is slightly different than the first equation: here we are using $x_\\text{d}$ instead of $x_\\text{e}$ and $u_\\text{d}$ instead of $u_\\text{e}$.  This is because we want our controller to track a desired trajectory $(x_\\text{d}(t), u_\\text{d}(t))$ rather than just stabilize the equilibrium point $(x_\\text{e}, u_\\text{e})$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "control, pvtol_closed = ct.create_statefbk_iosystem(pvtol, K)\n",
    "print(control, \"\\n\")\n",
    "print(pvtol_closed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This command will usually generate a warning saying that python control \"cannot verify system output is system state\".  This happens because we specified an output function `pvtol_output` when we created the system model, and python-control does not have a way of checking that the output function returns the entire state (which is needed if we are going to do full-state feedback).\n",
    "\n",
    "This warning could be avoided by passing the argument `None` for the system output function, in which case python-control returns the full state as the output (and it knows that the full state is being returned as the output)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bedcb0c0"
   },
   "source": [
    "## Closed loop system simulation\n",
    "\n",
    "For this simple example, we set the target for the system to be a \"step\" input that moves the system 1 meter to the right.\n",
    "\n",
    "We start by defining a short function to visualize the output using a collection of plots:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility function to plot the results in a useful way\n",
    "def plot_results(t, x, u, fig=None):\n",
    "    # Set the size of the figure\n",
    "    if fig is None:\n",
    "        fig = plt.figure(figsize=(10, 6))\n",
    "\n",
    "    # Top plot: xy trajectory\n",
    "    plt.subplot(2, 1, 1)\n",
    "    lines = plt.plot(x[0], x[1])\n",
    "    plt.xlabel('x [m]')\n",
    "    plt.ylabel('y [m]')\n",
    "    plt.axis('equal')\n",
    "\n",
    "    # Mark starting and ending points\n",
    "    color = lines[0].get_color()\n",
    "    plt.plot(x[0, 0], x[1, 0], 'o', color=color, fillstyle='none')\n",
    "    plt.plot(x[0, -1], x[1, -1], 'o', color=color, fillstyle='full')\n",
    "\n",
    "\n",
    "    # Time traces of the state and input\n",
    "    plt.subplot(2, 4, 5)\n",
    "    plt.plot(t, x[1])\n",
    "    plt.xlabel('Time t [sec]')\n",
    "    plt.ylabel('y [m]')\n",
    "\n",
    "    plt.subplot(2, 4, 6)\n",
    "    plt.plot(t, x[2])\n",
    "    plt.xlabel('Time t [sec]')\n",
    "    plt.ylabel('theta [rad]')\n",
    "\n",
    "    plt.subplot(2, 4, 7)\n",
    "    plt.plot(t, u[0])\n",
    "    plt.xlabel('Time t [sec]')\n",
    "    plt.ylabel('$F_1$ [N]')\n",
    "\n",
    "    plt.subplot(2, 4, 8)\n",
    "    plt.plot(t, u[1])\n",
    "    plt.xlabel('Time t [sec]')\n",
    "    plt.ylabel('$F_2$ [N]')\n",
    "    plt.tight_layout()\n",
    "\n",
    "    return fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we generate a step response and plot the results.  Because our closed loop system takes as inputs $x_\\text{d}$ and $u_\\text{d}$, we need to set those variable to values that would correspond to our step input.  In this case, we are taking a step in the $x$ coordinate, so we set $x_\\text{d}$ to be $1$ in that coordinate starting at $t = 0$ and continuing for some sufficiently long period of time ($15$ seconds):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate a step response by setting xd, ud\n",
    "Tf = 15\n",
    "T = np.linspace(0, Tf, 100)\n",
    "xd = np.outer(np.array([1, 0, 0, 0, 0, 0]), np.ones_like(T))\n",
    "ud = np.outer(ueq, np.ones_like(T))\n",
    "ref = np.vstack([xd, ud])\n",
    "\n",
    "response = ct.input_output_response(pvtol_closed, T, ref, xeq)\n",
    "fig = plot_results(response.time, response.states, response.outputs[6:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f014e660"
   },
   "source": [
    "This controller does a pretty good job.  We see in the top plot the $x$, $y$ projection of the trajectory, with the open circle indicating the starting point and the closed circle indicating the final point.  The bottom set of plots show the altitude and pitch as functions of time, as well as the input forces.  All of the signals look reasonable.\n",
    "\n",
    "The limitations of the linear controller can be seen if we take a larger step, say 10 meters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "xd = np.outer(np.array([10, 0, 0, 0, 0, 0]), np.ones_like(T))\n",
    "ref = np.vstack([xd, ud])\n",
    "response = ct.input_output_response(pvtol_closed, T, ref, xeq)\n",
    "fig = plot_results(response.time, response.states, response.outputs[6:])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4luxppVpm6Xo"
   },
   "source": [
    "We now see that the trajectory looses significant altitude ($> 2.5$ meters).  This is because the linear controller sees a large initial error and so it applies very large input forces to correct for the error ($F_1 \\approx -10$ N at $t = 0$.  This causes the aircraft to pitch over to a large angle (almost $-60$ degrees) and this causes a large loss in altitude.\n",
    "\n",
    "We will see in the [Lecture 6](cds110-L6a_kincar-trajgen) how to remedy this problem by making use of feasible trajectory generation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
